import numpy as np
import math
from astropy.io import fits

# Combine and transform Planck Stokes IQU to galactic coordinates
def Convert_Planck(StokesI, StokesQ, StokesU):
    # Creating numpy array for the total intensity
    NewI = np.copy(StokesI[0].data)
    # Converting the Stokes Q and U parameters to the new frame
    NewQ, NewU, NewIp, AngleO, AngleB = Convert_StokesQU(StokesQ[0].data, StokesU[0].data)
    # Creating the polarization fraction array
    NewP = 100.0*NewIp/NewI
    # Converting the units for the intensity arrays from Kcmb to mJy/arcsec^2
    # see Planck 2013 results. IX. HFI spectral response
    NewI = Convert_Kcmb_MJy_sr(NewI)
    NewQ = Convert_Kcmb_MJy_sr(NewQ)
    NewU = Convert_Kcmb_MJy_sr(NewU)
    NewIp = Convert_Kcmb_MJy_sr(NewIp)
    AngleO = 180.0/math.pi * AngleO
    AngleB = 180.0/math.pi * AngleB
    # Creating HDUList object to be returned
    newfile = fits.HDUList()
    # Creating the new Headers
    HeaderI = Fix_Header(StokesI[0].header, 'STOKES I', 'MJy/sr')
    HeaderQ = Fix_Header(StokesI[0].header, 'STOKES Q', 'MJy/sr')
    HeaderU = Fix_Header(StokesI[0].header, 'STOKES U', 'MJy/sr')
    HeaderIp = Fix_Header(StokesI[0].header, 'POL FLUX', 'MJy/sr')
    HeaderP = Fix_Header(StokesI[0].header, 'PERCENT POL', 'percent ')
    HeaderO = Fix_Header(StokesI[0].header, 'POL ANGLE', 'deg     ')
    HeaderB = Fix_Header(StokesI[0].header, 'ROTATED POL ANGLE', 'deg     ')
    # Populating the HDUList
    newfile.append(fits.ImageHDU(data=NewI, header=HeaderI))
    newfile.append(fits.ImageHDU(data=NewQ, header=HeaderQ))
    newfile.append(fits.ImageHDU(data=NewU, header=HeaderU))
    newfile.append(fits.ImageHDU(data=NewIp, header=HeaderIp))
    newfile.append(fits.ImageHDU(data=NewP, header=HeaderP))
    newfile.append(fits.ImageHDU(data=AngleO, header=HeaderO))
    newfile.append(fits.ImageHDU(data=AngleB, header=HeaderB))
    # Returning FITS container
    return newfile

# Function to convert units from Kcmb to MJy/sr
def Convert_Kcmb_MJy_sr(Array):
    # see Planck 2013 results. IX. HFI spectral response
    # Conversion factor from Kcmb to MJy per steradian (see Table 6 at 353 GHz)
    C = 246.543
    # Converting the array
    New_Array = C * Array
    # Returning the resulting converted array
    return New_Array

# Function to convert Stokes Q and U arrays from Planck standard to Galactic coordinates
def Convert_StokesQU(OldQ, OldU):
    # Creating the polarized intensity array
    Ip = (OldQ**2.0 + OldU**2.0)**0.5
    # Relation originally from Juan Soler's calculations
    # https://github.com/solerjuan/magnetar/tree/master
    psi=0.5*np.arctan2(OldU,OldQ)	# in radians
    ex=-np.sin(psi)
    ey=np.cos(psi)
    AngleP = np.arctan2(ex,ey)
    AngleB = np.arctan2(ey,-ex)
    # Finding the new stokes parameters
    NewQ = Ip*np.cos(2.0*AngleP)
    NewU = Ip*np.sin(2.0*AngleP)
    # Returning the corrected Stokes parameters
    return NewQ, NewU, Ip, AngleP, AngleB

# Modifying headers
def Fix_Header(Headerfile, ExtName, Units):
    # Keywords to transform: TTYPE1, TUNIT1, BUNIT
    # Creating new header by copying the original
    NewHeader = Headerfile.copy()
    # Fixing the cards for the header
    NewHeader['EXTNAME'] = (ExtName, 'ID of the HDU')
    NewHeader['TTYPE1'] = ExtName
    NewHeader['TUNIT1'] = Units
    NewHeader['BUNIT'] = Units
    # Returning new header
    return NewHeader